package com.geotmt.billingcenter.config.datasources;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.DruidDataSourceFactory;
import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.SqlSessionTemplate;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.lang.Nullable;
import org.springframework.transaction.PlatformTransactionManager;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Description 多数据源配置类
 * @Author      yanghanwei
 * @Mail        yanghanwei@geotmt.com
 * @Date        2019/3/19 14:07
 * @Version     v1
 **/
@Configuration
public class DataSourceConfiguration implements ApplicationContextAware, InitializingBean {
    private ApplicationContext context;
    @Nullable
    private Map<Object, Object> targetDataSources = new HashMap<>();
    @Nullable
    private Object defaultTargetDataSource;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        if (context == null) {
            context = applicationContext;
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        MatchProperties properties = context.getBean(MatchProperties.class);
        for (Properties prop : properties.getSource()) {
            String clazz = prop.getProperty("type");
            String label = prop.getProperty("label");
            DataSource dataSource = null;
            if (DruidDataSource.class.getName().equals(clazz)) {
                dataSource = DruidDataSourceFactory.createDataSource(prop);
            }
            if (dataSource != null) {
                if (label.equals(DataSourceContextHolder.MYSQL)) {
                    this.defaultTargetDataSource = dataSource;
                } else {
                    this.targetDataSources.put(label, dataSource);
                }
            }
        }
    }

    @Bean(name = "sqlSessionFactory")
    public SqlSessionFactory sqlSessionFactory() throws Exception {
        SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
        sqlSessionFactoryBean.setDataSource(roundRobinDataSourceProxy());
        return sqlSessionFactoryBean.getObject();
    }

    @Bean(name = "roundRobinDataSourceProxy")
    public DataSource roundRobinDataSourceProxy() {
        final int readSize = targetDataSources.size();
        AbstractRoutingDataSource proxy = new AbstractRoutingDataSource() {
            private AtomicInteger count = new AtomicInteger(0);
            @Override
            protected Object determineCurrentLookupKey() {
                String lookupKey = DataSourceContextHolder.get();
                if (lookupKey == null) {
                    //默认数据源，解决ShiroRealm注入的service事务失效或注入的dao无数据源问题
                    return null;
                }
                if (DataSourceContextHolder.HIVE.equals(lookupKey)) {
                    return lookupKey;
                }
                int number = count.getAndAdd(1);
                int idx = number % readSize;
                System.out.println("DataSource >> " + DataSourceContextHolder.MYSQL + (idx + 1));
                return DataSourceContextHolder.MYSQL + (idx + 1);
            }
        };

        proxy.setDefaultTargetDataSource(defaultTargetDataSource);
        proxy.setTargetDataSources(targetDataSources);
        return proxy;
    }

    @Bean
    public SqlSessionTemplate sqlSessionTemplate(SqlSessionFactory sqlSessionFactory) {
        return new SqlSessionTemplate(sqlSessionFactory);
    }

    @Bean
    public PlatformTransactionManager annotationDrivenTransactionManager(@Qualifier(value = "roundRobinDataSourceProxy") DataSource roundRobinDataSourceProxy) {
        return new DataSourceTransactionManager(roundRobinDataSourceProxy);
    }

}
